#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2021/4/22 13:08
# @Author  : LiShan
# @Email   : lishan_1997@126.com
# @File    : Gpu-Test.py
# @Note    : this is note
import os

os.environ["CUDA_VISIBLE_DEVICES"] = '0,1,2,3'
import torch
from torch import nn
import torchvision

print(torch.__version__)
print(torchvision.__version__)

print("cuda:", torch.cuda.is_available())  # GPU是否可用
print("Devices:%d" % torch.cuda.device_count())  # GPU个数
print("Current_Devices:%d" % torch.cuda.current_device())  # 当前GPU索引号
print(torch.cuda.get_device_name(0))  # 输出GPU名称
print([torch.cuda.get_device_properties(i) for i in range(torch.cuda.device_count())])  # GPU详细信息

device = torch.device('cuda' if torch.cuda.is_available else 'cpu')

# 在CPU上的张量计算
x = torch.tensor([1, 2, 3])
y = x ** 2
print(y, y.device)

# 在GPU上的张量计算
x = torch.tensor([1, 2, 3]).to(device)
y = x ** 2
print(y, y.device)

# z = y + x.cpu()  # 报错，不同设备间的数据不能直接计算

# 在CPU上的模型
module = nn.Linear(3, 1)
print(list(module.parameters())[0], list(module.parameters())[0].device)  # device(type='cpu')

# 在GPU上的模型
module = nn.Linear(3, 1).to(device)
print(list(module.parameters())[0], list(module.parameters())[0].device)  # device=(type='cuda', index=0)
